"""
快速测试脚本

验证智能会议安排系统的所有组件是否正常工作
"""

import sys
import os

def test_imports():
    """测试1: 检查所有模块是否能正常导入"""
    print("="*80)
    print("测试1: 模块导入")
    print("="*80)

    try:
        from agentic_system.environments.env_package.meeting_scheduling import (
            MeetingSchedulingEnv,
            MeetingSchedulingMultiProcessEnv,
            meeting_projection,
            meeting_memory_projection,
            MeetingReward
        )
        print("✓ 会议调度环境模块导入成功")
    except Exception as e:
        print(f"✗ 会议调度环境模块导入失败: {e}")
        return False

    try:
        from agentic_system.environments.env_package.meeting_scheduling.env_manager import (
            MeetingSchedulingEnvManager,
            create_meeting_env_manager
        )
        print("✓ 环境管理器模块导入成功")
    except Exception as e:
        print(f"✗ 环境管理器模块导入失败: {e}")
        return False

    try:
        from agentic_system.environments.prompts.meeting_prompts import (
            build_full_prompt,
            add_observation_to_messages,
            SYSTEM_PROMPT
        )
        print("✓ 提示词模块导入成功")
    except Exception as e:
        print(f"✗ 提示词模块导入失败: {e}")
        return False

    print("\n所有模块导入成功！\n")
    return True


def test_single_env():
    """测试2: 单环境功能"""
    print("="*80)
    print("测试2: 单环境功能")
    print("="*80)

    try:
        from agentic_system.environments.env_package.meeting_scheduling import MeetingSchedulingEnv

        env = MeetingSchedulingEnv(max_steps=20, seed=42)
        obs, info = env.reset()

        print(f"✓ 环境创建成功")
        print(f"  会议主题: {info['meeting_topic']}")
        print(f"  必选人员: {info['required_count']}")
        print(f"  可选人员: {info['optional_count']}")

        # 测试查询动作
        action = {
            'action_type': 'query_calendar',
            'parameters': {
                'persons': env.all_people[:2]
            }
        }
        obs, reward, done, info = env.step(action, valid=1)
        print(f"✓ 查询动作执行成功，奖励: {reward}")

        return True

    except Exception as e:
        print(f"✗ 单环境测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_projection():
    """测试3: 投影函数"""
    print("\n" + "="*80)
    print("测试3: 投影函数")
    print("="*80)

    try:
        from agentic_system.environments.env_package.meeting_scheduling import meeting_projection

        test_action = '''
        查询日程
        ```json
        {
            "action_type": "query_calendar",
            "parameters": {
                "persons": ["李明", "王芳"]
            }
        }
        ```
        '''

        action_dicts, valids = meeting_projection([test_action])

        if valids[0] == 1 and action_dicts[0]['action_type'] == 'query_calendar':
            print(f"✓ 投影函数解析成功")
            print(f"  动作类型: {action_dicts[0]['action_type']}")
            print(f"  参数: {action_dicts[0]['parameters']}")
            return True
        else:
            print(f"✗ 投影函数解析失败")
            return False

    except Exception as e:
        print(f"✗ 投影函数测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_reward_calculator():
    """测试4: 奖励计算"""
    print("\n" + "="*80)
    print("测试4: 奖励计算")
    print("="*80)

    try:
        from agentic_system.environments.env_package.meeting_scheduling import MeetingReward

        calculator = MeetingReward()

        # 测试理想情况
        scores = calculator.calculate_reward(
            required_available=3,
            required_total=3,
            optional_available=3,
            optional_total=3,
            historical_match=True,
            negotiations_count=0,
            steps_count=4,
            max_steps=20
        )

        total_reward = scores['total_reward']

        if total_reward > 0:
            print(f"✓ 奖励计算成功")
            print(f"  总奖励: {total_reward:.2f}")
            print(f"  明细: {scores}")
            return True
        else:
            print(f"✗ 奖励计算异常")
            return False

    except Exception as e:
        print(f"✗ 奖励计算测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_multiprocess_env():
    """测试5: 矢量化环境"""
    print("\n" + "="*80)
    print("测试5: 矢量化环境")
    print("="*80)

    try:
        from agentic_system.environments.env_package.meeting_scheduling import MeetingSchedulingMultiProcessEnv

        env = MeetingSchedulingMultiProcessEnv(
            seed=42,
            env_num=2,
            max_steps=20
        )

        observations, infos = env.reset()

        print(f"✓ 矢量化环境创建成功，环境数: {len(env)}")

        # 测试批量步骤
        text_actions = [
            '''
            ```json
            {"action_type": "query_calendar", "parameters": {"persons": ["李明"]}}
            ```
            ''',
            '''
            ```json
            {"action_type": "query_history", "parameters": {}}
            ```
            '''
        ]

        observations, rewards, dones, infos = env.step(text_actions, valids=[1, 1])

        print(f"✓ 批量步骤执行成功")
        print(f"  奖励: {rewards}")

        env.close()
        return True

    except Exception as e:
        print(f"✗ 矢量化环境测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_env_manager():
    """测试6: 环境管理器"""
    print("\n" + "="*80)
    print("测试6: 环境管理器")
    print("="*80)

    try:
        from agentic_system.environments.env_package.meeting_scheduling.env_manager import create_meeting_env_manager

        config = {
            'seed': 42,
            'env_num': 2,
            'max_steps': 20,
            'is_train': True,
            'use_memory': False
        }

        env_manager = create_meeting_env_manager(config)

        print(f"✓ 环境管理器创建成功")
        print(f"  环境名称: {env_manager.env_name}")

        # 测试重置
        observations = env_manager.reset()

        if 'text' in observations and len(observations['text']) == 2:
            print(f"✓ 环境管理器重置成功")
            return True
        else:
            print(f"✗ 环境管理器观察格式错误")
            return False

    except Exception as e:
        print(f"✗ 环境管理器测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_prompts():
    """测试7: 提示词模板"""
    print("\n" + "="*80)
    print("测试7: 提示词模板")
    print("="*80)

    try:
        from agentic_system.environments.prompts.meeting_prompts import build_full_prompt

        user_request = "下周找个时间开产品规划会，必须参加的有李明、王芳。"

        messages = build_full_prompt(user_request, use_memory=False)

        if len(messages) == 2 and messages[0]['role'] == 'system' and messages[1]['role'] == 'user':
            print(f"✓ 提示词模板构建成功")
            print(f"  消息数: {len(messages)}")
            return True
        else:
            print(f"✗ 提示词模板格式错误")
            return False

    except Exception as e:
        print(f"✗ 提示词模板测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def run_all_tests():
    """运行所有测试"""
    print("\n")
    print("#"*80)
    print("#" + " "*20 + "智能会议安排系统 - 快速测试" + " "*25 + "#")
    print("#"*80)
    print("\n")

    results = []

    # 运行测试
    results.append(("模块导入", test_imports()))
    results.append(("单环境功能", test_single_env()))
    results.append(("投影函数", test_projection()))
    results.append(("奖励计算", test_reward_calculator()))
    results.append(("矢量化环境", test_multiprocess_env()))
    results.append(("环境管理器", test_env_manager()))
    results.append(("提示词模板", test_prompts()))

    # 统计结果
    print("\n\n" + "="*80)
    print("测试结果汇总")
    print("="*80)

    total_tests = len(results)
    passed_tests = sum(1 for _, result in results if result)

    for name, result in results:
        status = "✓ 通过" if result else "✗ 失败"
        print(f"{name:20s} {status}")

    print("-"*80)
    print(f"总计: {passed_tests}/{total_tests} 通过")

    if passed_tests == total_tests:
        print("\n🎉 所有测试通过！系统运行正常。")
        return True
    else:
        print(f"\n⚠️  有 {total_tests - passed_tests} 个测试失败，请检查错误信息。")
        return False


if __name__ == '__main__':
    success = run_all_tests()
    sys.exit(0 if success else 1)
